package com.strp.mptcp.reassembler;

import com.strp.mptcp.ACK.AckHeader;
import com.strp.mptcp.ACK.manager.AckManager;
import com.strp.mptcp.common.DataChunk;
import com.strp.mptcp.ACK.SelectiveAckHeader;

import java.io.File;
import java.io.RandomAccessFile;
import java.nio.channels.FileChannel;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.ReentrantLock;

/**
 * Reassembler：负责数据重组和重传管理
 * 主要功能：
 * 1. 按序重组数据块
 * 2. 管理乱序数据缓存
 * 3. 实现快速重传和超时重传
 * 4. 支持滑动窗口流量控制
 * 5. 提供数据完整性校验
 *
 * @author 曹行阳
 * @date 2025/6/12 11:34
 * @version 0.1
 */
public class Reassembler {
    private static final long CHUNK_SIZE = 1024L * 64;          // 64KB
    private static final long RETRANSMIT_TIMEOUT_MS = 500L;     // 重传超时
    private static final int MAX_DUP_ACKS = 3;                  // 快速重传阈值
    private static final int WINDOW_SIZE = 1000;                // 接收窗口大小
    private static final int MAX_OUT_OF_ORDER = 1000;           // 最大乱序缓存数
    private static final long CLEANUP_INTERVAL_MS = 1000;       // 清理间隔

    private final FileChannel fileChannel;
    private final SortedMap<Long, DataChunk> outOfOrderBuffer = new TreeMap<>();
    private final Map<Long, Integer> dupAckCount = new ConcurrentHashMap<>();
    private final Set<Long> retransmitRequested = ConcurrentHashMap.newKeySet();
    private final AtomicLong nextExpectedSeq = new AtomicLong(0);
    private final AtomicLong windowStart = new AtomicLong(0);
    private final AtomicLong windowEnd = new AtomicLong(WINDOW_SIZE);
    private Long gapSince = null;

    private final ReentrantLock lock = new ReentrantLock();
    private final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(2);
    private final BlockingQueue<DataChunk> writeQueue = new LinkedBlockingQueue<>();
    private volatile boolean running = true;

    public Reassembler(File targetFile) throws Exception {
        this.fileChannel = new RandomAccessFile(targetFile, "rw").getChannel();

        // 启动重传定时任务
        scheduler.scheduleAtFixedRate(this::checkTimeouts,
                RETRANSMIT_TIMEOUT_MS,
                RETRANSMIT_TIMEOUT_MS,
                TimeUnit.MILLISECONDS);

        // 启动清理任务
        scheduler.scheduleAtFixedRate(this::cleanup,
                CLEANUP_INTERVAL_MS,
                CLEANUP_INTERVAL_MS,
                TimeUnit.MILLISECONDS);

        // 启动写入线程
        startWriterThread();
    }

    private void startWriterThread() {
        Thread writerThread = new Thread(() -> {
            while (running) {
                try {
                    DataChunk chunk = writeQueue.poll(100, TimeUnit.MILLISECONDS);
                    if (chunk != null) {
                        writeToDisk(chunk);
                    }
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    break;
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        });
        writerThread.setDaemon(true);
        writerThread.start();
    }

    public void processChunk(DataChunk chunk) {
        long seq = chunk.header().sequence();

        // 检查是否在接收窗口内
        if (seq < windowStart.get() || seq >= windowEnd.get()) {
            // 发送 SACK 通知发送方调整窗口
            sendWindowUpdate();
            return;
        }

        lock.lock();
        try {
            if (seq < nextExpectedSeq.get()) {
                // 重复块，增加重复计数
                int count = dupAckCount.getOrDefault(seq, 0) + 1;
                dupAckCount.put(seq, count);

                // 检查是否需要快速重传
                if (count >= MAX_DUP_ACKS && !retransmitRequested.contains(seq)) {
                    requestFastRetransmit(seq);
                }
                return;
            }

            if (seq == nextExpectedSeq.get()) {
                // 正常按序
                writeToDisk(chunk);
                nextExpectedSeq.incrementAndGet();

                // 处理缓存的乱序数据
                processBuffer();

                // 更新窗口
                updateWindow();

                // 发送 ACK
                sendAck();
            } else {
                // 乱序数据，缓存
                if (outOfOrderBuffer.size() < MAX_OUT_OF_ORDER) {
                    outOfOrderBuffer.put(seq, chunk);
                    if (gapSince == null) {
                        gapSince = System.currentTimeMillis();
                    }
                    // 发送 SACK
                    sendSelectiveAck();
                } else {
                    // 缓存已满，请求重传
                    requestRetransmission(nextExpectedSeq.get());
                }
            }
        } finally {
            lock.unlock();
        }
    }

    private void processBuffer() {
        while (true) {
            DataChunk next = outOfOrderBuffer.get(nextExpectedSeq.get());
            if (next == null) break;

            outOfOrderBuffer.remove(nextExpectedSeq.get());
            writeToDisk(next);
            nextExpectedSeq.incrementAndGet();
        }
    }

    private void writeToDisk(DataChunk chunk) {
        try {
            writeQueue.put(chunk);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }

    private void doWriteToDisk(DataChunk chunk) {
        try {
            long position = chunk.header().sequence() * CHUNK_SIZE;
            fileChannel.position(position);
            fileChannel.write(chunk.payload().nioBuffer());
            chunk.payload().release();
        } catch (Exception e) {
            e.printStackTrace();
            // 写入失败，请求重传
            requestRetransmission(chunk.header().sequence());
        }
    }

    private void updateWindow() {
        long newStart = nextExpectedSeq.get();
        long newEnd = newStart + WINDOW_SIZE;
        windowStart.set(newStart);
        windowEnd.set(newEnd);
    }

    private void sendAck() {
        AckManager.sendAck(new AckHeader(nextExpectedSeq.get() - 1));
    }

    private void sendSelectiveAck() {
        Set<Long> receivedSeqs = new HashSet<>(outOfOrderBuffer.keySet());
        receivedSeqs.add(nextExpectedSeq.get() - 1);
        AckManager.sendSelectiveAck(new SelectiveAckHeader(nextExpectedSeq.get(), receivedSeqs));
    }

    private void sendWindowUpdate() {
        // 发送窗口更新通知
        AckManager.sendWindowUpdate(windowStart.get(), windowEnd.get());
    }

    private void requestFastRetransmit(long seq) {
        retransmitRequested.add(seq);
        AckManager.requestRetransmission(seq);
    }

    private void requestRetransmission(long seq) {
        if (!retransmitRequested.contains(seq)) {
            retransmitRequested.add(seq);
            AckManager.requestRetransmission(seq);
        }
    }

    private void checkTimeouts() {
        lock.lock();
        try {
            if (gapSince != null) {
                long now = System.currentTimeMillis();
                if (now - gapSince >= RETRANSMIT_TIMEOUT_MS) {
                    requestRetransmission(nextExpectedSeq.get());
                    gapSince = now;
                }
            }
        } finally {
            lock.unlock();
        }
    }

    private void cleanup() {
        // 清理过期的重传请求
        long now = System.currentTimeMillis();
        retransmitRequested.removeIf(seq ->
                now - gapSince >= RETRANSMIT_TIMEOUT_MS * 3);

        // 清理重复 ACK 计数
        dupAckCount.clear();
    }

    public void close() throws Exception {
        running = false;
        scheduler.shutdownNow();
        if (!scheduler.awaitTermination(5, TimeUnit.SECONDS)) {
            scheduler.shutdownNow();
        }
        fileChannel.close();
    }
}
